# faiss_rag/embed.py（完整替换，不用保留原来的代码）
import ollama
import numpy as np
from typing import List

# 全局初始化Ollama客户端（连接本地Ollama服务）
client = ollama.Client()

def get_embedding(text: str) -> List[float]:
    """用Ollama的all-minilm模型生成文本嵌入向量"""
    try:
        # 调用Ollama嵌入API，生成向量
        response = client.embeddings(
            model="all-minilm:latest",  # 刚才拉取的嵌入模型
            prompt=text  # 要生成嵌入的文本
        )
        # 返回向量（转为list，确保和原来的格式一致）
        return response["embedding"]
    except Exception as e:
        print(f"Ollama嵌入生成失败：{e}")
        raise  # 抛出错误，方便定位问题

# 测试用（可选，运行embed.py可验证）
if __name__ == "__main__":
    vec = get_embedding("测试文本")
    print(f"向量维度：{len(vec)}")  # 应该输出384，和VECTOR_DIM匹配